from faker import Faker
import random
from datetime import datetime
from dateutil.relativedelta import relativedelta
from data_mock.utils import FileUtil, MysqlUtils_AI

fake = Faker('zh_CN')

# ========== 基础配置 ==========
HOSPITAL_CODE = 'H001'
HOSPITAL_NAME = '山西省肿瘤医院'
# CANCERS = ['肺癌', '胃癌']
cancer = '肺癌'
# VISIT_TYPES = ['门诊', '住院']
visit_type = '住院'
PATIENT_TYPES = {
    '10034': 'A1',
    '10033': 'B17'
}


# ========== 时间范围配置 ==========
def generate_month_range(start_month: str, end_month: str):
    """生成月份范围"""
    start = datetime.strptime(start_month, "%Y-%m")
    end = datetime.strptime(end_month, "%Y-%m")
    months = []
    current = start
    while current <= end:
        months.append(current.strftime("%Y-%m"))
        current += relativedelta(months=1)
    return months


PROJECT_MONTHS = generate_month_range("2025-07", "2025-08")


# ========== 核心数据生成函数 ==========
def generate_patient_type_detail_records():
    sql_statements = []
    record_count = 0

    for month in PROJECT_MONTHS:
        for i in range(100):
            for code, p_type in PATIENT_TYPES.items():
                record_count += 1

                patient_type_count = random.randint(10, 100)

                # 把月份字符串转成 datetime 对象
                start_dt = datetime.strptime(f"{month}-01 00:00:00", "%Y-%m-%d %H:%M:%S")
                end_dt = datetime.strptime(f"{month}-28 23:59:59", "%Y-%m-%d %H:%M:%S")

                create_time = fake.date_time_between(
                    start_date=start_dt,
                    end_date=end_dt
                ).strftime('%Y-%m-%d %H:%M:%S')

                data = {
                    'hospital_code': HOSPITAL_CODE,
                    'hospital_name': HOSPITAL_NAME,
                    'cancer': cancer,
                    'visit_type': visit_type,
                    'project_month': month,
                    'patient_type_code': code,
                    'patient_type': p_type,
                    'patient_type_count': patient_type_count,
                    'create_time': create_time
                }

                sql = _generate_sql('patient_type_detail', data)
                sql_statements.append(sql)

    return sql_statements


# ========== SQL 生成函数 ==========
def _generate_sql(table, data):
    columns = ', '.join([f'`{k}`' for k in data.keys()])
    values = []
    for v in data.values():
        if v is None:
            values.append('NULL')
        elif isinstance(v, (int, float)):
            values.append(str(v))
        else:
            escaped_value = str(v).replace("'", "''")
            values.append(f"'{escaped_value}'")
    return f"INSERT INTO `{table}` ({columns}) VALUES ({', '.join(values)});"


# ========== 执行生成 ==========
if __name__ == "__main__":
    records = generate_patient_type_detail_records()
    print(records)
    # 写入数据库
    # MysqlUtils_AI.insert_data_to_hub(records, 'patient_type_detail')
    # 或写入 SQL 文件
    # FileUtil.generate_sql_file(records, "patient_type_detail.sql")
